#!/usr/bin/env python

import rospy
import os
import os.path as osp
from sensor_msgs.msg import PointCloud2, Image
from sensor_msgs import point_cloud2
import cv_bridge
import rospkg
import pcl
import cv2
import numpy as np


base_path = None
for path in rospkg.get_ros_package_path().split(':'):
    if 'yunfei' in path:
        base_path = path
assert base_path, "base path not set correctly"
data_path = osp.join(base_path, 'bag_dealer', "bags", "data")

index1 = 1
index2 = 1

def point_cloud_saver_cb(msg):
    global index1
    gen = point_cloud2.read_points(msg, field_names=('x', 'y', 'z', 'intensity'))
    data = np.zeros(0)
    flag = True
    for idx, p in enumerate(gen):
        if flag:
            data = np.array([p[0], p[1], p[2], p[3]], dtype=np.float32).reshape(-1, 4)
            flag = False
        else:
            data = np.vstack((data, np.array([p[0], p[1], p[2], p[3]], dtype=np.float32)))
    pc = pcl.PointCloud_PointXYZI(data)
    pcl.save(pc, osp.join(data_path, 'pcd', str(index1) + '.pcd'))
    index1 += 1


def image_saver_cb(msg):
    global index2
    bridge = cv_bridge.CvBridge()
    img = bridge.imgmsg_to_cv2(msg)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    cv2.imwrite(osp.join(data_path, 'img', str(index2) + '.jpg'), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
    index2 += 1

def data_finer():
    rospy.init_node('data_finer', anonymous=True)
    rospy.Subscriber("feature_extraction/experimental_region", PointCloud2, point_cloud_saver_cb)
    # rospy.Subscriber("/usb_cam0/image_raw", Image, image_saver_cb)
    rospy.spin()

if __name__ == "__main__":
    data_finer()